package com.nd.bdc.presto.service;

import com.nd.bdc.presto.connection.PrestoJdbcClient;
import com.nd.bdc.presto.entity.CatalogsView.*;
import com.nd.bdc.presto.entity.CatalogsView;
import com.nd.bdc.presto.entity.PrestoSelectResult;
import com.nd.bdc.presto.util.JdbcUtil;
import com.nd.bdc.presto.util.SqlSelectManager;
import com.nd.bdc.presto.util.ThreadManager;
import io.prestosql.jdbc.PrestoConnection;
import io.prestosql.jdbc.PrestoStatement;
import io.prestosql.sql.parser.StatementSplitter;
import org.springframework.stereotype.Service;
import javax.annotation.Resource;
import java.sql.Connection;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.atomic.AtomicReference;
import java.util.stream.Collectors;
import static io.prestosql.sql.parser.StatementSplitter.isEmptyStatement;

@Service
public class PrestoServiceImpl implements PrestoService{

    public static final int DEFAULT_SQL_RESULTSET_FETCH_SIZE = 100;

    @Resource
    PrestoJdbcClient prestoJdbcClient;

    @Resource
    SqlSelectManager sqlSelectManager;

    @Resource
    PrestoWebSocketService prestoWebSocketService;


    @Override
    public PrestoSelectResult query(String catalog, String schema, String sql, Integer startRow, Integer fetchSize) throws Throwable{
        if (startRow == null) {
            startRow = 1;
        }
        if (fetchSize == null) {
            fetchSize = DEFAULT_SQL_RESULTSET_FETCH_SIZE;
        }

        if (fetchSize < 1) {
            fetchSize = 1;
        }
        if (fetchSize > 1000) {
            fetchSize = 1000;
        }

        final int startRowFinal = startRow;
        final int fetchSizeFinal = fetchSize;
        PrestoSelectResult result = sqlSelectManager.select(prestoJdbcClient.getSchemaConnection(catalog,schema),sql,startRowFinal,fetchSizeFinal);

        return result;
    }

    @Override
    public void execute(String channel, String catalog, String schema, String sql, Integer resultsetFetchSize){
        List<String> sqlList = new ArrayList<>();
        try{

            if(sql != null && sql.isEmpty()){
                return;
            }

            StatementSplitter splitter = new StatementSplitter(sql);
            for (StatementSplitter.Statement split : splitter.getCompleteStatements()) {
                if (!isEmptyStatement(split.statement())) {
                    sqlList.add(split.statement());
                }
            }

            if (!isEmptyStatement(splitter.getPartialStatement())) {
                sqlList.add(splitter.getPartialStatement());
            }
        }catch (Throwable t){
            prestoWebSocketService.sendExceptionMessage(channel,t,"parseSqlException",true);
            return;
        }


        Connection cn = null;
        Statement st = null;

        // 1.向前端发送任务开始消息
        prestoWebSocketService.sendStartMessage(channel, sqlList.size());

        try {
            cn = prestoJdbcClient.getSchemaConnection(catalog, schema);
            st = cn.createStatement();
            if(resultsetFetchSize != null){
                st.setFetchSize(resultsetFetchSize);
            }
        } catch (Throwable t) {
            // 2.向前端发送任务异常信息
            prestoWebSocketService.sendExceptionMessage(channel,t,"executionConnectionException",true);
            // 3.向前端发送任务结束信息
            prestoWebSocketService.sendFinishMessage(channel);
            return;
        }

        long startTime = System.currentTimeMillis();
        int totalCount = sqlList.size();
        PrestoWebSocketService.SQLExecutionStat sqlExecutionStat = new PrestoWebSocketService.SQLExecutionStat(totalCount);

        for(int i=0;i<totalCount;i++){
            String sqlStatment = sqlList.get(i);
            try
            {
                execute(channel,sqlExecutionStat, cn, st,sqlStatment , i);
                sqlExecutionStat.increaseSuccessCount();

            }
            catch (SQLException e)
            {
                sqlExecutionStat.increaseExceptionCount();
                prestoWebSocketService.sendSqlExceptionMessage(channel,sqlStatment,i,e,"executionSQLException");
            }
        }

        JdbcUtil.closeStatement(st);
        JdbcUtil.closeConnection(cn);

        sqlExecutionStat.setTaskDuration(System.currentTimeMillis() - startTime);

        // 向前端发送任务执行完成信息
        prestoWebSocketService.sendFinishMessage(channel, sqlExecutionStat);

    }


    @Override
    public List<String> catalogs() throws Throwable {
        String catalogQuerySql = "show catalogs";
        List<String> catalogList = query(null,null,catalogQuerySql,null,null)
                .getRows().stream().map( e -> e.get("Catalog").toString() ).collect(Collectors.toList());
        return catalogList;
    }

    @Override
    public List<String> schemas(String catalog) throws Throwable {
        String schemasQuery = "show schemas";
        List<String> schemaList = query( catalog, null, schemasQuery, null, null)
                .getRows().stream().map(e -> e.get("Schema").toString()).collect(Collectors.toList());
        return schemaList;
    }

    @Override
    public List<String> tables(String catalog,String schema) throws Throwable {
        String tableQuery = "show tables";
        List<String> tableList = query( catalog, schema, tableQuery, null, null)
                .getRows().stream().map(e -> e.get("Table").toString()).collect(Collectors.toList());
        return tableList;
    }

    @Override
    public List<String> columns(String catalog, String schema, String table) throws Throwable {
        String columnQuery = String.format("show columns from %s", table);
        List<String> columnList = query( catalog, schema, columnQuery, null, null)
                .getRows().stream().map(
//                        e -> String.format("%s (%s)",e.get("Column").toString(),e.get("Type").toString())
                        e -> e.get("Column").toString()
                ).collect(Collectors.toList());
        return columnList;
    }

    @Override
    public Object kill(String queryId) throws Throwable {
        Object res = prestoJdbcClient.killQuery(queryId);
        return res;
    }

    @Override
    public CatalogsView catalog() throws Throwable {
        // 初始化catalogs
        CatalogsView catalogsView = new CatalogsView();
        String catalogQuerySql = "show catalogs";
        List<String> catalogList = query(null,null,catalogQuerySql,null,null)
                .getRows().stream().map( e -> e.get("Catalog").toString() ).collect(Collectors.toList());
        for (String catalog : catalogList) {

            // 初始化catalog
            CatalogView catalogView = new CatalogView(catalog);
            // 添加catalog到catalogs
            catalogsView.getCatalogs().add(catalogView);

            String schemasQuery = "show schemas";
            List<String> schemaList = query( catalog, null, schemasQuery, null, null)
                    .getRows().stream().map(e -> e.get("Schema").toString()).collect(Collectors.toList());

            for (String schema : schemaList) {
                String tableQuery = "show tables";
                List<String> tableList = query( catalog, schema, tableQuery, null, null)
                        .getRows().stream().map(e -> e.get("Table").toString()).collect(Collectors.toList());
                // 初始化schema
                CatalogView.SchemaView schemaView = new CatalogView.SchemaView(schema);
                // 添加tableList 到schema
                schemaView.setTables(tableList);
                // 添加schemaView到schemas
                catalogView.getSchemas().add(schemaView);
            }
        }

        return catalogsView;
    }

    protected void execute(String channel,PrestoWebSocketService.SQLExecutionStat sqlExecutionStat,
                           Connection cn, Statement st, String sqlStatement, int sqlStatementIndex) throws SQLException
    {
        long startTime = System.currentTimeMillis();

        String sql = sqlStatement;

        PrestoConnection conn = cn.unwrap(PrestoConnection.class);


        // 添加任务进度监控
        AtomicReference<String> queryId = new AtomicReference<>();
        st.unwrap(PrestoStatement.class).setProgressMonitor(queryStats -> {
            //设置监听器（可选），可监听presto任务执行状况
            queryId.set(queryStats.getQueryId());    //获取presto任务ID（可用该ID终止任务）
            //向前端发送任务执行进度信息
            if(queryStats.getProgressPercentage().isPresent()){
                ThreadManager.getThreadPollProxy().execute(
                        () -> {
                            prestoWebSocketService.sendProgressMonitorMessage(channel,sqlStatement,sqlStatementIndex,queryStats);
                        }
                );
            }

        });

        boolean isResultSet = st.execute(sql);

        sqlExecutionStat.increaseSqlDuration(System.currentTimeMillis() - startTime);

        //关闭进度监控
        st.unwrap(PrestoStatement.class).clearProgressMonitor();

        // 查询操作
        if (isResultSet)
        {
            ResultSet rs = st.getResultSet();

            PrestoSelectResult sqlSelectResult = sqlSelectManager.select(sql, rs, 1,
                    DEFAULT_SQL_RESULTSET_FETCH_SIZE);

            //向前端发送查询结果信息
            prestoWebSocketService.sendSqlSuccessMessage(channel,sqlStatement,sqlStatementIndex,sqlSelectResult);
        }
        else
        {
            int updateCount = st.getUpdateCount();
            // 更新操作
            if (updateCount > -1)
            {
                // 向前端发送更新成功信息
                prestoWebSocketService.sendSqlSuccessMessage(channel,sqlStatement,sqlStatementIndex,updateCount);
            }
            // 其他操作
            else
            {
                // 向前端发送成功信息
                prestoWebSocketService.sendSqlSuccessMessage(channel,sqlStatement,sqlStatementIndex);
            }
        }
    }

}
